package com.shlh.saas.es;

import com.alibaba.fastjson.JSON;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.lang.math.NumberUtils;
import org.apache.http.HttpHost;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.ClearScrollResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.sort.ScoreSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.stereotype.Component;

import java.util.*;

@Component
@Slf4j
public class EsRestClient {

    public static Map<String, RestHighLevelClient> clientMap = new HashMap<>();

    @Resource
    private EsConfigProperties esConfigProperties;

    private static final RequestOptions COMMON_OPTIONS;

    static {
        RequestOptions.Builder builder = RequestOptions.DEFAULT.toBuilder();
        COMMON_OPTIONS = builder.build();
    }

    @PostConstruct
    public void initialize() {
        List<EsClusterConfig> esConfigs = esConfigProperties.getEsConfigs();
        for (EsClusterConfig esConfig : esConfigs) {
            log.info("initializing.config.name: {},nodes: {} ",esConfig.getName(), esConfig.getNodes());
            RestHighLevelClient restHighLevelClient = initRestClient(esConfig);
            if (restHighLevelClient != null) {
                clientMap.put(esConfig.getName(), restHighLevelClient);
            }else {
                log.error("initializing.config.name: {},nodes: {} failed",esConfig.getName(), esConfig.getNodes());
            }
        }
    }


    private RestHighLevelClient initRestClient(EsClusterConfig esClusterConfig) {
        String[] ipPortArray = esClusterConfig.getNodes().split(",");
        List<HttpHost> httpHostList = new ArrayList<>(ipPortArray.length);
        for (String ipPort : ipPortArray) {
            String[] ipPortInfo = ipPort.split(":");
            if (ipPortInfo.length == 2) {
                HttpHost httpHost = new HttpHost(ipPortInfo[0], NumberUtils.toInt(ipPortInfo[1]));
                httpHostList.add(httpHost);
            }
        }
        HttpHost[] httpHosts = new HttpHost[httpHostList.size()];
        httpHostList.toArray(httpHosts);

        RestClientBuilder builder = RestClient.builder(httpHosts)
            // 设置连接超时时间，默认为30秒，增加到120秒
            .setRequestConfigCallback(requestConfigBuilder -> 
                requestConfigBuilder.setConnectTimeout(120000)  // 连接超时
                .setSocketTimeout(120000))  // Socket超时
            // 设置连接池
            .setHttpClientConfigCallback(httpClientBuilder -> 
                httpClientBuilder.setMaxConnTotal(200)  // 最大连接数
                .setMaxConnPerRoute(100));  // 每个路由最大连接数
            
        RestHighLevelClient restHighLevelClient = new RestHighLevelClient(builder);
        log.info("已创建ES客户端，集群: {}, 节点: {}, 连接超时: 120秒", esClusterConfig.getName(), esClusterConfig.getNodes());
        return restHighLevelClient;
    }

    private static RestHighLevelClient getClient(String clusterName) {
        return clientMap.get(clusterName);
    }

    /**
     * 新增文档
     * @param esIndexInfo
     * @param esSourceData
     * @return
     */
    public static boolean insertDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData){
        try {
            IndexRequest indexRequest = new IndexRequest(esIndexInfo.getIndexName());
            indexRequest.source(esSourceData.getData());
            indexRequest.id(esSourceData.getDocId());
            getClient(esIndexInfo.getClusterName()).index(indexRequest, COMMON_OPTIONS);
            return true;
        }catch (Exception e){
            log.error("insertDoc.exception: {}",e.getMessage(),e);
        }
        return false;
    }

    /**
     * 更新文档
     * @param esIndexInfo
     * @param esSourceData
     * @return
     */
    public static boolean updateDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData){
        try {
            UpdateRequest updateRequest = new UpdateRequest();
            updateRequest.index(esIndexInfo.getIndexName());
            updateRequest.id(esSourceData.getDocId());
            updateRequest.doc(esSourceData.getData());

            getClient(esIndexInfo.getClusterName()).update(updateRequest, COMMON_OPTIONS);
            return true;
        }catch (Exception e){
            log.error("updateDoc.exception: {}",e.getMessage(),e);
        }
        return false;
    }

    /**
     * 批量更新文档
     * @param esIndexInfo
     * @param esSourceDataList
     * @return
     */
    public static boolean batchUpdateDoc(EsIndexInfo esIndexInfo,
                                         List<EsSourceData> esSourceDataList){
        try {
            boolean flag = false;
            BulkRequest bulkRequest = new BulkRequest();
            for (EsSourceData esSourceData : esSourceDataList) {
                String docId = esSourceData.getDocId();
                if(StringUtils.isBlank(docId)){
                    UpdateRequest updateRequest = new UpdateRequest();
                    updateRequest.index(esIndexInfo.getIndexName());
                    updateRequest.id(esSourceData.getDocId());
                    updateRequest.doc(esSourceData.getData());
                    bulkRequest.add(updateRequest);
                    flag = true;
                }
            }

            if (flag){
                BulkResponse bulkResponse = getClient(esIndexInfo.getClusterName()).bulk(bulkRequest, COMMON_OPTIONS);
                if (bulkResponse.hasFailures()){
                    return false;
                }
            }

            return true;
        }catch (Exception e){
            log.error("batchUpdateDoc.exception: {}",e.getMessage(),e);
        }
        return false;
    }

    /**
     * 删除文档
     *
     */
    public static boolean delete(EsIndexInfo esIndexInfo){
        try {
            DeleteByQueryRequest deleteByQueryRequest =
                    new DeleteByQueryRequest(esIndexInfo.getIndexName());
            deleteByQueryRequest.setQuery(QueryBuilders.matchAllQuery());
            BulkByScrollResponse response = getClient(
                            esIndexInfo.getClusterName()).deleteByQuery(
                                    deleteByQueryRequest, COMMON_OPTIONS
            );
            long deleted = response.getDeleted();
            log.info("deleteDoc.deleted: {}", deleted);
            return true;
        }catch (Exception e){
            log.error("deleteDoc.exception: {}",e.getMessage(),e);
        }
        return false;
    }

    /**
     * 根据id删除文档
     *
     */
    public static boolean delete(EsIndexInfo esIndexInfo, String docId){
        try {
            DeleteRequest deleteRequest = new DeleteRequest(esIndexInfo.getIndexName());
            deleteRequest.id(docId);
            DeleteResponse response = getClient(esIndexInfo.getClusterName()).delete(deleteRequest, COMMON_OPTIONS);
            log.info("deleteDoc.response: {}", JSON.toJSONString(response));
            return true;
        }catch (Exception e){
            log.error("deleteDoc.exception: {}",e.getMessage(),e);
        }
        return false;
    }

    /**
     * 文档是否存在
     */
    public static boolean isExistsDocById(EsIndexInfo esIndexInfo, String docId){
        try {
            GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
            getRequest.id(docId);
            return getClient(esIndexInfo.getClusterName()).exists(getRequest, COMMON_OPTIONS);
        }catch (Exception e){
            log.error("isExistsDocById.exception: {}",e.getMessage(),e);
        }
        return false;
    }

    /**
     * 根据id查询文档
     */
    public static Map<String, Object> queryById(EsIndexInfo esIndexInfo, String docId) {
        try {
            GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
            getRequest.id(docId);
            GetResponse response = getClient(esIndexInfo.getClusterName()).get(getRequest, COMMON_OPTIONS);
            Map<String, Object> source = response.getSource();
            return source;

        } catch (Exception e) {
            log.error("queryById.exception: {}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 根据id获取指定字段
     */
    public static Map<String, Object> queryById(EsIndexInfo esIndexInfo, String docId,
                                                String[] fields) {
        try {
            GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
            getRequest.id(docId);
            FetchSourceContext fetchSourceContext = new FetchSourceContext(true, fields, null);
            getRequest.fetchSourceContext(fetchSourceContext);
            GetResponse response = getClient(esIndexInfo.getClusterName()).get(getRequest, COMMON_OPTIONS);
            Map<String, Object> source = response.getSource();
            return source;

        } catch (Exception e) {
            log.error("queryById.exception: {}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 搜索文档
     */
    public static SearchResponse searchWithTermQuery(EsIndexInfo esIndexInfo, EsSearchRequest esSearchRequest) {
        try {
            BoolQueryBuilder bq = esSearchRequest.getBq();
            String[] fields = esSearchRequest.getFields();
            Long from = esSearchRequest.getFrom();
            Long size = esSearchRequest.getSize();
            Long minutes = esSearchRequest.getMinustes();
            String scrollTime = esSearchRequest.getScrollTime(); // 获取scrollTime参数
            Boolean needScroll = esSearchRequest.getNeedScroll();
            String sortName = esSearchRequest.getSortName();
            String sortOrder = esSearchRequest.getSortOrder();
            boolean trackTotalHits = esSearchRequest.isTrackTotalHits(); // 获取trackTotalHits参数

            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchSourceBuilder.query(bq);
            searchSourceBuilder.fetchSource(fields, null).from(Math.toIntExact(from)).size(Math.toIntExact(size));

            int searchTimeout = 120000;
            searchSourceBuilder.timeout(TimeValue.timeValueMillis(searchTimeout));

            // 设置trackTotalHits参数，确保返回精确计数
            if (trackTotalHits) {
                searchSourceBuilder.trackTotalHits(true);
                log.debug("ES查询启用trackTotalHits=true，确保返回精确计数");
            }

            if (Objects.nonNull(esSearchRequest.getHighlightBuilder())){
                searchSourceBuilder.highlighter(esSearchRequest.getHighlightBuilder());
            }

            if (StringUtils.isNotBlank(sortName)){
                searchSourceBuilder.sort(sortName);
            }

            searchSourceBuilder.sort(new ScoreSortBuilder().order(SortOrder.DESC));
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.searchType(SearchType.DEFAULT);
            searchRequest.indices(esIndexInfo.getIndexName());
            searchRequest.source(searchSourceBuilder);
            
            if (needScroll){
                // 优先使用scrollTime参数，如果没有则使用minutes参数兼容旧代码
                if (StringUtils.isNotBlank(scrollTime)) {
                    Scroll scroll = new Scroll(TimeValue.parseTimeValue(scrollTime, "scroll_time"));
                    searchRequest.scroll(scroll);
                    log.debug("使用scrollTime参数创建Scroll: {}", scrollTime);
                } else if (minutes != null) {
                Scroll scroll = new Scroll(TimeValue.timeValueMinutes(minutes));
                searchRequest.scroll(scroll);
                    log.debug("使用minutes参数创建Scroll: {}分钟", minutes);
                } else {
                    // 默认5分钟
                    Scroll scroll = new Scroll(TimeValue.timeValueMinutes(5));
                    searchRequest.scroll(scroll);
                    log.debug("未指定Scroll时间，使用默认值: 5分钟");
                }
            }

            // 记录完整的ES查询参数（包括trackTotalHits）
            log.debug("ES查询完整参数: index={}, query={}, from={}, size={}, trackTotalHits={}",
                     esIndexInfo.getIndexName(), bq.toString(), from, size, trackTotalHits);

            SearchResponse search = getClient(esIndexInfo.getClusterName()).search(searchRequest, COMMON_OPTIONS);
            return search;

        }catch (Exception e){
            log.error("searchWithTermQuery.exception: {}",e.getMessage(),e);
        }
        return null;
    }

    /**
     * 执行scroll查询，获取下一批结果
     *
     * @param scrollId scroll ID
     * @param scrollTime scroll保持时间
     * @return 搜索响应
     */
    public static SearchResponse searchScroll(String scrollId, String scrollTime) {
        try {
            log.debug("执行scroll查询: scrollId={}, scrollTime={}", scrollId, scrollTime);

            // 创建scroll查询请求
            SearchScrollRequest searchScrollRequest =
                new SearchScrollRequest(scrollId);
            searchScrollRequest.scroll(TimeValue.parseTimeValue(scrollTime, "scroll_time"));

            // 从第一个可用的客户端执行查询
            SearchResponse searchResponse = null;
            for (Map.Entry<String, RestHighLevelClient> entry : clientMap.entrySet()) {
                try {
                    searchResponse = entry.getValue().scroll(searchScrollRequest, COMMON_OPTIONS);
                    break;  // 成功获取结果，退出循环
                } catch (Exception e) {
                    log.warn("从ES集群{}执行scroll查询失败: {}", entry.getKey(), e.getMessage());
                    // 继续尝试下一个客户端
                }
            }

            if (searchResponse == null) {
                log.error("所有ES集群执行scroll查询失败");
                return null;
            }

            log.debug("scroll查询成功: 获取到{}条结果", 
                    searchResponse.getHits() == null ? 0 : searchResponse.getHits().getHits().length);
            return searchResponse;

        } catch (Exception e) {
            log.error("执行scroll查询失败: scrollId={}, error={}", scrollId, e.getMessage(), e);
            return null;
        }
    }

    /**
     * 清理scroll上下文，释放资源
     *
     * @param scrollId scroll ID
     * @return 是否成功
     */
    public static boolean clearScroll(String scrollId) {
        try {
            log.debug("清理scroll上下文: scrollId={}", scrollId);

            // 创建清理scroll请求
            ClearScrollRequest clearScrollRequest = 
                new ClearScrollRequest();
            clearScrollRequest.addScrollId(scrollId);

            // 从第一个可用的客户端执行清理
            boolean success = false;
            for (Map.Entry<String, RestHighLevelClient> entry : clientMap.entrySet()) {
                try {
                    ClearScrollResponse clearScrollResponse = 
                        entry.getValue().clearScroll(clearScrollRequest, COMMON_OPTIONS);
                    success = clearScrollResponse.isSucceeded();
                    break;  // 成功执行，退出循环
                } catch (Exception e) {
                    log.warn("从ES集群{}清理scroll上下文失败: {}", entry.getKey(), e.getMessage());
                    // 继续尝试下一个客户端
                }
            }

            if (!success) {
                log.warn("所有ES集群清理scroll上下文失败");
                return false;
            }

            log.debug("成功清理scroll上下文");
            return true;

        } catch (Exception e) {
            log.error("清理scroll上下文失败: scrollId={}, error={}", scrollId, e.getMessage(), e);
            return false;
        }
    }

    /**
     * 执行scroll查询获取所有匹配文档
     * @param esIndexInfo 索引信息
     * @param esSearchRequest 搜索请求
     * @return 所有匹配文档的列表
     */
    public static List<Map<String, Object>> scrollSearchAll(EsIndexInfo esIndexInfo, EsSearchRequest esSearchRequest) {
        List<Map<String, Object>> allResults = new ArrayList<>();
        String scrollId = null;

        try {
            // 初始搜索请求
            SearchResponse searchResponse = searchWithTermQuery(esIndexInfo, esSearchRequest);
            if (searchResponse == null) {
                return allResults;
            }

            scrollId = searchResponse.getScrollId();
            SearchHits searchHits = searchResponse.getHits();

            // 处理第一批结果
            processSearchHits(searchHits, allResults);

            // 继续滚动查询直到没有更多数据
            while (searchHits != null && searchHits.getHits() != null && searchHits.getHits().length > 0) {
                // 使用scrollId继续获取下一批数据
                SearchResponse scrollResponse = searchScroll(scrollId, esSearchRequest.getScrollTime());

                if (scrollResponse == null) {
                    break;
                }

                // 更新scrollId
                scrollId = scrollResponse.getScrollId();
                searchHits = scrollResponse.getHits();

                // 处理当前批次结果
                processSearchHits(searchHits, allResults);

                log.debug("已获取{}条数据，继续滚动查询...", allResults.size());
            }

            log.info("滚动查询完成，共获取{}条数据", allResults.size());
        } catch (Exception e) {
            log.error("滚动查询异常", e);
        } finally {
            // 清理scroll上下文
            if (scrollId != null) {
                clearScroll(scrollId);
            }
        }

        return allResults;
    }

    /**
     * 处理搜索命中的结果
     */
    private static void processSearchHits(SearchHits searchHits, List<Map<String, Object>> results) {
        if (searchHits == null || searchHits.getHits() == null) {
            return;
        }

        for (SearchHit hit : searchHits.getHits()) {
            results.add(hit.getSourceAsMap());
        }
    }

}
